import os
from typing import Optional, Dict, Any, Union
import pytorch_lightning as pl
from pytorch_lightning.loggers.logger import Logger
from pytorch_lightning.utilities.rank_zero import rank_zero_only
from argparse import Namespace


class SwanLabLogger(Logger):
    """
    自定义 SwanLab 日志记录器
    """

    def __init__(
        self,
        project: Optional[str] = None,
        workspace: Optional[str] = "Wizard",  # 已经注册过了
        experiment_name: Optional[str] = None,
        description: Optional[str] = None,
        logdir: Optional[str] = None,
    ):
        super().__init__()
        # 延迟导入 swanlab，避免在不需要时加载
        try:
            import swanlab

            self._swanlab = swanlab
        except ImportError:
            raise ImportError(
                "You want to use `swanlab` logger which is not installed yet, "
                "install it with `pip install swanlab`."
            )

        self._project = project
        self._workspace = workspace
        self._experiment_name = experiment_name
        self._description = description
        self._logdir = logdir
        self._experiment = None

    @property
    def name(self) -> str:
        """返回实验名称"""
        return self._experiment_name or "swanlab_experiment"

    @property
    def version(self) -> str:
        """返回实验版本"""
        return "0.1.0"

    @property
    @rank_zero_only
    def experiment(self):
        """返回 swanlab 实验实例"""
        if self._experiment is None:
            self._experiment = self._swanlab.init(
                project=self._project,
                workspace=self._workspace,
                name=self._experiment_name,
                description=self._description,
                logdir=self._logdir,
            )
        return self._experiment

    @rank_zero_only
    def log_metrics(
        self, metrics: Dict[str, float], step: Optional[int] = None
    ) -> None:
        """记录指标"""
        assert rank_zero_only.rank == 0, "experiment tried to log from global_rank != 0"

        # 移除不需要的指标
        metrics.pop("epoch", None)

        # 记录指标到 swanlab
        self.experiment.log(metrics, step=step)

    @rank_zero_only
    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
        """记录超参数"""
        # 转换参数为字典格式
        if isinstance(params, Namespace):
            params = vars(params)
        
        # 如果参数是嵌套字典，需要展平
        def flatten_dict(d, parent_key='', sep='.'):
            items = []
            for k, v in d.items():
                new_key = f"{parent_key}{sep}{k}" if parent_key else k
                if isinstance(v, dict):
                    items.extend(flatten_dict(v, new_key, sep=sep).items())
                else:
                    items.append((new_key, v))
            return dict(items)
        
        # 展平参数字典
        if isinstance(params, dict):
            params = flatten_dict(params)
        
        # 更新实验配置
        self.experiment.config.update(params)

    @rank_zero_only
    def finalize(self, status: str) -> None:
        """结束实验"""
        if self._experiment is not None:
            self._experiment.finish()
        self._experiment = None
